from ultralytics.data.base import BaseDataset
from ultralytics.data.augment import Compose, RandomHSV, RandomFlip
from ultralytics.utils.instance import Instances
from pathlib import Path
import glob
import os
import cv2
import numpy as np
from ultralytics.utils import DEFAULT_CFG

class MyDataset(BaseDataset):
    def __init__(self, img_path, **kwargs):
        super().__init__(img_path, **kwargs)

    def get_labels(self):
        f = []  # image files
        for p in self.img_path if isinstance(self.img_path, list) else [self.img_path]:
            p = Path(p)  # os-agnostic
            if p.is_dir():  # dir
                f += glob.glob(str(p / "**" / "*.*"), recursive=True)
                # F = list(p.rglob('*.*'))  # pathlib
            elif p.is_file():  # file
                with open(p) as t:
                    t = t.read().strip().splitlines()
                    parent = str(p.parent) + os.sep
                    f += [x.replace("./", parent) if x.startswith("./") else x for x in t]  # local to global path
                    # F += [p.parent / x.lstrip(os.sep) for x in t]  # local to global path (pathlib)
            else:
                raise FileNotFoundError(f"{self.prefix}{p} does not exist")
        txt_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in {"txt"})
        im_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in {"jpg"})

        labels = []
        for i in range(len(txt_files)):
            lab_dict = dict()
            txtfile = txt_files[i]
            imgfile = im_files[i]
            lab_dict["im_file"] = imgfile

            img = cv2.imread(imgfile)
            height, width, _ = img.shape
            lab_dict["shape"] = (height, width)

            with open(txtfile, 'r') as f:
                label_txt = f.readlines()
            cls = []
            bboxes = []
            for line in label_txt:
                cls_box = line[:-1]
                cls_box_list = cls_box.split()
                cls.append(float(cls_box_list[0]))
                bboxes.append([float(cls_box_list[1]), float(cls_box_list[2]), float(cls_box_list[3]), float(cls_box_list[4])])

            lab_dict["cls"] = cls
            lab_dict["bboxes"] = bboxes
            lab_dict["normalized"] = True
            lab_dict["bbox_format"] = "xyxy"
            labels.append(lab_dict)
        #print(labels)
        return labels

    def build_transforms(self, hyp=None):
        if self.augment:
            return Compose([RandomFlip(), RandomHSV()])

mydataset = MyDataset(img_path='./coco8', augment=True)

# load_image 加载图片到ram测试
# 打印输入信息
#print('ims:\n', mydataset.ims)
#print('im_files:\n', mydataset.im_files)
#print('npy_files:\n', mydataset.npy_files)
# 单张图片加载测试
ims, im_hw0, im_hw = mydataset.load_image(1, rect_mode=True)
save1dir = './load_img.jpg'
cv2.imwrite(save1dir, ims)
#print('ims:\n', [i if i is None else 'load success' for i in mydataset.ims])

# cache_image_to_disk 磁盘保存测试
#mydataset.cache_images_to_disk(1)

# transforms 图像增强测试
print(DEFAULT_CFG)

labels = mydataset.get_image_and_label(1)
bboxes = np.array(labels["bboxes"])
bbox_format = labels["bbox_format"]
normalized = labels["normalized"]
img = labels['img']
im_file = labels["im_file"]
cls = labels['cls']
new_labels = {'im_file':im_file, 'cls':cls, 'img':img,
              'instances':Instances(bboxes=bboxes, bbox_format=bbox_format, normalized=normalized)
              }

trans_img = mydataset.transforms(new_labels)['img']

save2dir = './trans_img.jpg'
cv2.imwrite(save2dir, trans_img)

